package com.atguigu.aiproject.interceptor;

import com.atguigu.aiproject.constant.RedisKeyConstant;
import com.atguigu.aiproject.entity.User;
import com.atguigu.aiproject.utils.JwtUtils;
import io.jsonwebtoken.Claims;
import io.jsonwebtoken.ExpiredJwtException;
import jakarta.servlet.http.HttpServletRequest;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.data.redis.core.RedisTemplate;
import org.springframework.http.HttpStatus;
import org.springframework.http.server.ServerHttpRequest;
import org.springframework.http.server.ServerHttpResponse;
import org.springframework.http.server.ServletServerHttpRequest;
import org.springframework.stereotype.Component;
import org.springframework.web.socket.WebSocketHandler;
import org.springframework.web.socket.server.support.HttpSessionHandshakeInterceptor;
import reactor.util.annotation.NonNull;

import java.util.Map;
import java.util.concurrent.TimeUnit;

@Component
public class JwtHandshakeInterceptor extends HttpSessionHandshakeInterceptor {
    private final JwtUtils jwtUtils;
    private final RedisTemplate<String, Object> redisTemplate;
    private static final Logger log = LoggerFactory.getLogger(JwtHandshakeInterceptor.class);
    private static final String USER_JWT_KEY = RedisKeyConstant.USER_JWT_KEY;
    private static final String JWT_USER_KEY = RedisKeyConstant.JWT_USER_KEY;

    public JwtHandshakeInterceptor(JwtUtils jwtUtils, RedisTemplate<String, Object> redisTemplate) {
        this.jwtUtils = jwtUtils;
        this.redisTemplate = redisTemplate;
    }

    @Override
    public boolean beforeHandshake(@NonNull ServerHttpRequest request,
                                   @NonNull ServerHttpResponse response,
                                   @NonNull WebSocketHandler wsHandler,
                                   @NonNull Map<String, Object> attributes) throws Exception {

        // 检查是否是 WebSocket 握手请求
        if (request instanceof ServletServerHttpRequest servletRequest) {
            HttpServletRequest httpRequest = servletRequest.getServletRequest();

            // 从请求参数或头中获取 token
            String token = getTokenFromRequest(httpRequest);

            if (token == null || token.isEmpty()) {
                log.warn("WebSocket 握手失败: 未提供 token");
                response.setStatusCode(HttpStatus.UNAUTHORIZED);
                return false;
            }

            try {
                // 验证 token 并获取用户信息
                User user = validateTokenAndGetUser(token);
                if (user == null) {
                    log.warn("WebSocket 握手失败: token 验证失败");
                    response.setStatusCode(HttpStatus.UNAUTHORIZED);
                    return false;
                }

                // 将用户信息存入 attributes，在 WebSocketHandler 中可以通过 session.getAttributes() 获取
                attributes.put("userId", user.getId());
                attributes.put("username", user.getLoginName());
                attributes.put("user", user);

                log.info("WebSocket 握手成功: 用户 {}", user.getLoginName());
                return true;

            } catch (Exception e) {
                log.error("WebSocket 握手异常: {}", e.getMessage());
                response.setStatusCode(HttpStatus.UNAUTHORIZED);
                return false;
            }
        }

        return super.beforeHandshake(request, response, wsHandler, attributes);
    }

    /**
     * 从请求中获取 token
     * 支持从查询参数、请求头中获取
     */
    private String getTokenFromRequest(HttpServletRequest request) {
        // 1. 首先尝试从查询参数获取
        String token = request.getParameter("token");
        if (token != null && !token.isEmpty()) {
            return token;
        }

        // 2. 从 Authorization 头获取
        String authHeader = request.getHeader("Authorization");
        if (authHeader != null && authHeader.startsWith("Bearer ")) {
            return authHeader.substring(7);
        }

        // 3. 从 Sec-WebSocket-Protocol 头获取（某些 WebSocket 客户端使用）
        String protocolHeader = request.getHeader("Sec-WebSocket-Protocol");
        if (protocolHeader != null && protocolHeader.startsWith("Bearer")) {
            String[] protocols = protocolHeader.split(",");
            for (String protocol : protocols) {
                protocol = protocol.trim();
                if (protocol.startsWith("Bearer ")) {
                    return protocol.substring(7);
                }
            }
        }

        return null;
    }

    /**
     * 验证 token 并获取用户信息
     */
    private User validateTokenAndGetUser(String token) {
        try {
            // 1. 尝试从 Redis 获取用户信息
            User user = getUserFromRedis(token);
            if (user != null) {
                // 单端登录验证
                if (!isLatestToken(token, user.getId())) {
                    return null;
                }
                return user;
            }

            // 2. Redis 中没有，解析 JWT
            Claims claims = jwtUtils.extractAllClaims(token);

            // 3. 从 JWT 构建用户信息
            user = buildUserFromClaims(claims);

            // 单端登录验证
            if (!isLatestToken(token, user.getId())) {
                return null;
            }

            // 4. 缓存到 Redis
            safeCacheUser(token, user, claims.getExpiration().getTime());

            return user;

        } catch (ExpiredJwtException e) {
            safeDeleteFromRedis(token);
            log.warn("Token 已过期: {}", token);
            return null;
        } catch (Exception e) {
            safeDeleteFromRedis(token);
            log.error("Token 验证失败: {}", e.getMessage());
            return null;
        }
    }

    /**
     * 验证 token 是否有效（检查单端登录）
     */
    private boolean isLatestToken(String token, Integer userId) {
        try {
            String currentToken = (String) redisTemplate.opsForValue().get(USER_JWT_KEY + userId);
            if (currentToken == null) {
                return true;
            }
            return token.equals(currentToken);
        } catch (Exception e) {
            log.error("验证 token 失败", e);
            return true;
        }
    }

    /**
     * 从 Redis 获取用户信息
     */
    private User getUserFromRedis(String token) {
        try {
            String redisKey = JWT_USER_KEY + token;
            return (User) redisTemplate.opsForValue().get(redisKey);
        } catch (Exception e) {
            log.error("Redis 不可用，降级到 JWT 验证: {}", e.getMessage());
            return null;
        }
    }

    /**
     * 安全地删除 Redis 中的 token
     */
    private void safeDeleteFromRedis(String token) {
        try {
            redisTemplate.delete(JWT_USER_KEY + token);
        } catch (Exception e) {
            log.error("Redis 不可用，忽略删除: {}", e.getMessage());
        }
    }

    /**
     * 安全地缓存用户信息到 Redis
     */
    private void safeCacheUser(String token, User user, long expirationTime) {
        try {
            long expireTime = expirationTime - System.currentTimeMillis();
            if (expireTime > 0) {
                redisTemplate.opsForValue().set(JWT_USER_KEY + token, user, expireTime, TimeUnit.MILLISECONDS);
                redisTemplate.opsForValue().set(USER_JWT_KEY + user.getId(), token, expireTime, TimeUnit.MILLISECONDS);
            }
        } catch (Exception e) {
            log.error("Redis 不可用，忽略缓存: {}", e.getMessage());
        }
    }

    /**
     * 从 Claims 构建用户对象
     */
    private User buildUserFromClaims(Claims claims) {
        String username = (String) claims.get("username");
        Integer userId = (Integer) claims.get("userId");
        User user = new User();
        user.setLoginName(username);
        user.setId(userId);
        return user;
    }

    /**
     * 握手成功后的处理
     */
    @Override
    public void afterHandshake(@NonNull ServerHttpRequest request,
                               @NonNull ServerHttpResponse response,
                               @NonNull WebSocketHandler wsHandler,
                               Exception exception) {
        // 握手完成后的处理，可以在这里记录日志等
        if (exception == null) {
            log.debug("WebSocket 握手完成");
        } else {
            log.error("WebSocket 握手异常: {}", exception.getMessage());
        }
        super.afterHandshake(request, response, wsHandler, exception);
    }
}
